"""
管理向量数据库，进行文档存储和检索
向量存储
    1、初始化与集合管理：创建或加载Milvus向量数据库集合。
    2、文档向量化与存储：将分块后的文档转换为向量并存储。
    3、混合检索与重排序：结合稠密和稀疏向量进行检索，并通过重排序优化结果。
"""
import hashlib
from gc import enable

import torch.cuda
from langchain.docstore.document import Document
# 文本嵌入模型，将文本（如查询、文档片段）转化为向量（Embedding），供 Milvus 检索使用
from milvus_model.hybrid import BGEM3EmbeddingFunction
# AnnSearchRequest: 构建向量检索请求的对象，指定检索的集合、向量、TopK 等参数，是执行向量检索的核心载体。
# WeightedRanker: 混合检索（如 “向量检索 + 标量过滤”“多向量检索”）的结果排序器，通过加权策略融合多源检索结果，优化排序效果。
from pymilvus import AnnSearchRequest, DataType, MilvusClient, WeightedRanker
# CrossEncoder: 用于直接计算文本对相关性得分的模型，常用于文档重排序任务。
from sentence_transformers import CrossEncoder


from src.spark_edu_rag.rag_qa.core.document_processor import \
    document_split_chunk
from src.spark_edu_rag.utils import config_ini, get_logger, get_project_root
from src.spark_edu_rag.rag_qa.core.document_processor import document_split_chunk


class VectorStore:
    def __init__(self):
        self.logger = get_logger(__name__)
        self._config_loader = config_ini
        self.collection_name = self._config_loader.MILVUS.COLLECTION_NAME
        self.uri = f"http://{self._config_loader.MILVUS.HOST}:{self._config_loader.MILVUS.PORT}"
        self.client = MilvusClient(uri=self.uri, db_name=self._config_loader.MILVUS.DATABASE_NAME)

        self.model_m3_path = get_project_root.joinpath(self._config_loader.MODEL_PATH.BGE_M3_MODEL_PATH)
        self.embedding_function = BGEM3EmbeddingFunction(
            model_name=self.model_m3_path,
            use_fp16=self._config_loader.MODEL_PATH.M3_USE_FP16
        )
        self.m3_dense_dim = self.embedding_function.dim["dense"]

        self.device = "gpu" if torch.cuda.is_available() else "cpu"
        self.logger.info(f"使用设备: {self.device}")

        self.reranker = CrossEncoder(
            model_name_or_path=get_project_root.joinpath(self._config_loader.MODEL_PATH.RERANKER_MODEL_PATH),
            device=self.device
        )

    # 创建或者加载milvus集合
    def _create_or_load_collection(self):

        if not self.client.has_collection(self.collection_name):
            schema = self.client.create_schema(auto_id=False, enable_dynamic_field=True)
            schema.add_field(field_name="id", datatype=DataType.VARCHAR, is_primary=True, max_length=100)
            schema.add_field(field_name="text", datatype=DataType.VARCHAR, max_length=65535)
            # 稠密向量字段
            schema.add_field(field_name="dense_vector", datatype=DataType.FLOAT_VECTOR, dim=self.m3_dense_dim)
            # 稀疏向量字段
            schema.add_field(field_name="sparse_vector", datatype=DataType.SPARSE_FLOAT_VECTOR)
            schema.add_field(field_name="parent_id", datatype=DataType.VARCHAR, max_length=200)
            schema.add_field(field_name="parent_content", datatype=DataType.VARCHAR, max_length=65535)
            schema.add_field(field_name="source", datatype=DataType.VARCHAR, max_length=300)
            schema.add_field(field_name="create_time", datatype=DataType.VARCHAR, max_length=50)

            prepare_index_param = self.client.prepare_index_params()
            prepare_index_param.add_index(
                field_name="dense_vector",
                index_name="dense_vector_index",
                index_type="IVF_FLAT", # 索引类型：IVF_FLAT（倒排文件平坦索引）
                metric_type="IP", #  距离计算方式：内积（Inner Product）
                params={"nlist": 128},
            )
            prepare_index_param.add_index(
                field_name="sparse_vector",
                index_name="sparse_vector_index",
                index_type="SPARSE_INVERTED_INDEX", # 索引类型：SPARSE_INVERTED_INDEX（稀疏倒排索引）
                metric_type="IP", #  距离计算方式：内积（Inner Product）
                params={"drop_ratio_build": 0.2},
            )
            self.client.create_collection(
                collection_name=self.collection_name,
                schema=schema,
                index_params= prepare_index_param,
            )
            self.logger.info(f"创建集合: {self.collection_name}，创建的字段列表：{schema.fields}，索引参数：{prepare_index_param}")
        else:
            self.logger.info(f"集合已存在: {self.collection_name}")
        collection_status = self.client.get_collection_stats(self.collection_name)
        if collection_status and "row_count" in collection_status:
            self.logger.info(f"集合已加载: {self.collection_name}")
        else:
            self.client.load_collection(self.collection_name)
            self.logger.info(f"加载集合: {self.collection_name}")



    # 将子块文档转换为向量并插入milvus
    def add_document(self, documents: list[Document]):
        """
        将子块文档转换为向量并插入milvus
        :param documents: 子块文档列表
        :return:
        """
        texts = [doc.page_content for doc in documents]
        embeddings = self.embedding_function(texts)
        self.logger.info(f"嵌入向量：{embeddings}")

        # 存放保存到milvus实体数据
        entities = []
        for doc_index, doc in enumerate(documents):
            # 计算向量
            text_hash = hashlib.md5(doc.page_content.encode('utf-8')).hexdigest()
            sparse_vector = {}
            # 获取稀疏向量的非零索引和值
            row = embeddings["sparse"][doc_index]
            # 处理CSR格式稀疏矩阵
            coo_row = row.tocoo()
            row_index = coo_row.col  # 使用列索引作为键
            row_value = coo_row.data

            for entity_index, entity in zip(row_index, row_value):
                sparse_vector[entity_index] = entity
            entities.append({
                "id": text_hash,
                "text": doc.page_content,
                "dense_vector": embeddings["dense"][doc_index],
                "sparse_vector": sparse_vector,
                "parent_id": doc.metadata.get("parent_id"),
                "parent_content": doc.metadata.get("parent_content"),
                "source": doc.metadata.get("source"),
                "create_time": doc.metadata.get("create_time")
            })
        # # 插入到milvus
        if entities:
            self.client.upsert(
                collection_name=self.collection_name,
                data=entities
            )
            self.logger.info(f"插入 {len(entities)} 条实体到集合: {self.collection_name}")

    # 混合检索（稠密+稀疏） + 重排序 ，返回精准父文档
    def hybrid_search_with_rerank(self, query: str, cond: str, top_k: int) -> list[Document]:
        """
        混合检索（稠密+稀疏） + 重排序 ，返回精准父文档
        :param query: 查询文本，如：AI学科的课程内容是什么？
        :param cond: 查询条件，如：{"source": "AI"}
        :param top_k: 返回Top K个文档，默认从配置文件中读取
        :return: 精准父文档列表
        """
        # 计算查询向量
        query_embeddings = self.embedding_function([query])
        dense_query_vector = query_embeddings["dense"][0]
        sparse_query_vector = {}
        row = query_embeddings["sparse"][0]
        coo_row = row.tocoo()
        row_index = coo_row.col  # 使用列索引作为键
        row_value = coo_row.data
        # row_index = row.indices
        # row_value = row.data
        for entity_index, entity in zip(row_index, row_value):
            sparse_query_vector[entity_index] = entity
        # 构建过滤表达式，根据条件过滤文档来源
        # filter_expr = f"source == '{cond}'" if cond else ""
        filter_expr = f"source LIKE '%{cond}%'" if cond else ""

        # 构建AnnSearchRequest,构建稠密向量检索请求，设置过滤表达式
        dense_search_param = AnnSearchRequest(
            data = [dense_query_vector],
            anns_field = "dense_vector",  # 稀疏向量检索字段
            param = {"metric_type": "IP", "params": {"nprobe": 10}},  # 稀疏向量检索距离度量类型，IP表示点积
            limit =  top_k,
            expr = filter_expr
        )
        sparse_search_param = AnnSearchRequest(
            data = [sparse_query_vector],
            anns_field = "sparse_vector",
            param = {"metric_type": "IP", "params": {}},
            limit = top_k,
            expr = filter_expr
        )
        # 构建加权排序器，根据稠密向量距离（语义相似度）和稀疏向量（关键词匹配）加权排序
        ranker = WeightedRanker(1.0, 0.7)
        # result内容格式：[{'id': 1234567890123456789, 'parent_content': '父文档内容', 'source': 'AI', 'text': '子文档内容', 'dense_vector': [0.1, 0.2, ...], 'sparse_vector': {0: 0.5, 1: 0.3, ...}}]
        # [0]，原因是只有一个检索请求，返回结果为一个列表，每个元素为一个字典，包含检索结果
        results = self.client.hybrid_search(
            collection_name=self.collection_name,
            ranker=ranker,
            reqs=[dense_search_param, sparse_search_param],
            limit=top_k,
            output_fields=["parent_id", "parent_content", "source", "text", "dense_vector", "sparse_vector"]
        )[0]
        # 将检索结果，转换为langchain Document列表，包含父文档内容、来源、文本、稠密向量、稀疏向量
        child_chunks = [self._doc_from_hit(doc["entity"]) for doc in results]
        parent_docs = self._get_unique_parent_docs(child_chunks)

        # 重排序逻辑：父文档数量 < 2 时跳过排序，直接返回所有父文档
        if len(parent_docs) < 2:
            return parent_docs[:top_k]
        # 12. 父文档数量 >= 2时, 执行重排序(提升相关性)
        if parent_docs:
            # 12.1 构建"查询-文档"配对列表
            pairs = [[query, doc.page_content] for doc in parent_docs]
            # 12.2 计算相关性得分, 得分越高, 相关性越强
            scores = self.reranker.predict(pairs)
            # print(f'scores: {scores}')
            # 12.3 按得分降序排序
            ranked_parent_docs = [doc for _, doc in sorted(zip(scores, parent_docs), reverse=True)]
        else:
            # 13. 若父文档为空(无检索结果), 返回空列表.
            ranked_parent_docs = []

            # 14. 返回重拍后的Top-M个文档.
            return ranked_parent_docs[:top_k]

    # 将milvus检索结果，转换为langchain Document
    def convert_to_documents(self, results: list[dict]) -> list[Document]:
        """
        将milvus检索结果，转换为langchain Document
        :param results: milvus检索结果
        :return: langchain Document列表
        """
        documents = []
        for result in results:
            doc = Document(
                page_content=result["text"],
                metadata={
                    "parent_id": result["parent_id"],
                    "parent_content": result["parent_content"],
                    "source": result["source"],
                    "create_time": result["create_time"]
                }
            )
            documents.append(doc)
        return documents
    # 定义私有方法，从 Milvus 查询结果创建 Document 对象
    def _doc_from_hit(self, hit):
        # 创建并返回 Document 对象，填充内容和元数据
        return Document(
            page_content=hit.get("text"),
            metadata={
                "parent_id": hit.get("parent_id"),
                "parent_content": hit.get("parent_content"),
                "source": hit.get("source"),
                "create_time": hit.get("create_time")
            }
        )

    # 定义私有方法，从子块中提取去重的父文档
    def _get_unique_parent_docs(self, sub_chunks):
        # 初始化集合，用于存储已处理的父块内容（去重）
        parent_contents = set()
        # 初始化列表，用于存储唯一父文档
        unique_docs = []
        # 遍历所有子块
        for chunk in sub_chunks:
            # 获取子块的父块内容，默认为子块内容
            parent_content = chunk.metadata.get("parent_content", chunk.page_content)
            # 检查父块内容是否非空且未重复
            if parent_content and parent_content not in parent_contents:
                # 创建新的 Document 对象，包含父块内容和元数据
                unique_docs.append(Document(page_content=parent_content, metadata=chunk.metadata))
                # 将父块内容添加到去重集合
                parent_contents.add(parent_content)
        # 返回去重后的父文档列表
        return unique_docs


if __name__ == '__main__':
    vector_store = VectorStore()
    # print(f"连接Milvus: {vector_store.uri}, 集合: {vector_store.collection_name}")
    # print(vector_store.model_m3_path)
    # vector_store._create_or_load_collection()
    #
    #
    # documents = document_split_chunk(r"D:\WorkArea\WorkSpace\Python\spark-edu-rag\src\spark_edu_rag\resources\data\ai_data")
    # vector_store.add_document(documents)
    query = "AI学科的课程内容是什么"
    results = vector_store.hybrid_search_with_rerank(query=query, cond="ai", top_k=3)
    # 4. 打印检索结果
    print(f'results: {results}')        # 结果列表
    print(f'results: {len(results)}')   # 结果数量